import numpy as np
import matplotlib.pyplot as plt
from DOA import DOA


class ESPRIT(DOA):
    def __init__(self):
        super(ESPRIT, self).__init__()

    def esprit(self, array_signal):
        # 计算协方差矩阵
        R = np.dot(array_signal, array_signal.conj().T) / array_signal.shape[1]
        # 特征值分解并取得信号子空间
        D, U = np.linalg.eigh(R)  # D为特征值， U为特征向量
        Us = U[:, -self.num_signals:]  # 信号子空间
        # 角度估计
        Ux = Us[:-1, :]  # Us为M*M矩阵，前M-1构成Ux  Ux(sensors-1, signal)
        Uy = Us[1:, :]
        # 最小二乘返回(x, residuals, rank, s)，其中 x 是解向量
        Psi = np.linalg.lstsq(Ux, Uy, rcond=None)[0]  # 最小二乘解函数Ux * Phi = Uy
        T, Phi = np.linalg.eig(Psi)  # T是特征值，Phi是特征向量 phi=(signal,signal)
        # Theta = np.rad2deg(np.arccos(np.angle(np.diag(Phi))))
        # Theta = np.arccos(np.angle(np.diag(Phi)) * self.wavelength / (2 * np.pi * self.sensor_distance))  # 估计角度
        Theta = np.arccos(np.angle(T) * self.wavelength / (2 * np.pi * self.sensor_distance))  # 估计角度
        Theta = np.sort(Theta)
        return Theta



    def apply_esprit(self):
        array_signal = self.generate_signal()
        estimated_angles = self.esprit(array_signal)
        # 打印估计结果
        print("估计结果：", np.rad2deg(estimated_angles))

        # 绘制角度估计结果的直方图
        # plt.hist(estimated_angles, bins='auto', color='blue', alpha=0.7, rwidth=0.85)
        # plt.grid(axis='y', alpha=0.75)
        # plt.xlabel('角度')
        # plt.ylabel('频率')
        # plt.title('角度估计结果直方图')
        # plt.show()


esprit = ESPRIT()
esprit.apply_esprit()